from torch import nn

from model.conv_bn_relu import ConvBNRelu
from options import HiDDenConfiguration


class Decoder(nn.Module):
    def __init__(self,config:HiDDenConfiguration):
        super().__init__()
        self.channels = config.decoder_channels

        #print("message_length:", config.message_length)  # 确认是否与预期一致（如100）
        #print("decoder_channels:", self.channels)  # 确认是否与预期一致（如64）


        layers = [ConvBNRelu(3, self.channels)]

        for _ in range (config.decoder_blocks - 1):
            layers.append(ConvBNRelu(self.channels,self.channels))

        layers.append(ConvBNRelu(self.channels,config.message_length))
        #自适应平均池化层
        #将前面输出的特征图去掉 H和W维度
        layers.append(nn.AdaptiveAvgPool2d(output_size=(1, 1)))
        self.layers = nn.Sequential(*layers)

        #实现线性变换的核心模块
        #y = x * w +b  w为可以学习的权重矩阵，b为可以学习的偏执向量
        #用于对解码的信息进行微调
        self.linear = nn.Sequential(
            nn.Linear(config.message_length, config.message_length ),
            nn.Sigmoid()
        )


    def forward(self, image_with_wm):
        #print("输入含水印图像形状：", image_with_wm.shape)

        x = self.layers(image_with_wm)
        #print("卷积层+池化层输出形状：", x.shape)
        #删除两个维度
        x.squeeze_(3).squeeze_(2)
        #print("挤压后形状：", x.shape)
        x = self.linear(x)
        return x
